package com.jweb.framework.ws.client.security;

import java.io.IOException;
import java.security.Principal;
import java.security.cert.X509Certificate;
import java.util.Vector;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.UnsupportedCallbackException;
import org.apache.axiom.soap.SOAPEnvelope;
import org.apache.axiom.soap.SOAPFactory;
import org.apache.axiom.soap.SOAPMessage;
import org.apache.ws.security.WSSecurityEngine;
import org.apache.ws.security.WSSecurityEngineResult;
import org.apache.ws.security.WSSecurityException;
import org.apache.ws.security.WSUsernameTokenPrincipal;
import org.apache.ws.security.components.crypto.Crypto;
import org.apache.ws.security.handler.RequestData;
import org.apache.ws.security.handler.WSHandlerResult;
import org.apache.ws.security.message.token.Timestamp;
import org.apache.ws.security.util.WSSecurityUtil;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.ws.context.MessageContext;
import org.springframework.ws.soap.SoapMessage;
import org.springframework.ws.soap.axiom.support.AxiomUtils;
import org.springframework.ws.soap.saaj.SaajSoapMessage;
import org.springframework.ws.soap.security.AbstractWsSecurityInterceptor;
import org.springframework.ws.soap.security.WsSecuritySecurementException;
import org.springframework.ws.soap.security.WsSecurityValidationException;
import org.springframework.ws.soap.security.callback.CallbackHandlerChain;
import org.springframework.ws.soap.security.callback.CleanupCallback;
import org.springframework.ws.soap.security.wss4j.Wss4jSecuritySecurementException;
import org.springframework.ws.soap.security.wss4j.Wss4jSecurityValidationException;
import org.springframework.ws.soap.security.wss4j.callback.UsernameTokenPrincipalCallback;
import org.w3c.dom.Document;
/**
 * 修改人： Thomas
 * 修改内容：新增 
 * 类说明：
 */
public class Wss4jSecurityInterceptorJWeb extends AbstractWsSecurityInterceptor implements InitializingBean {
	
	public static final String SECUREMENT_USER_PROPERTY_NAME = "Wss4jSecurityInterceptor.securementUser";
    private int securementAction;
    private String securementActions;
    private Vector securementActionsVector;
    private String securementUsername;
    private CallbackHandler validationCallbackHandler;
    private int validationAction;
    private String validationActions;
    private Vector validationActionsVector;
    private String validationActor;
    private Crypto validationDecryptionCrypto;
    private Crypto validationSignatureCrypto;
    private boolean timestampStrict;
    private boolean enableSignatureConfirmation;
    private int validationTimeToLive;
    private int securementTimeToLive;
    private final Wss4jHandlerJWeb handler = new Wss4jHandlerJWeb();
    private final WSSecurityEngine securityEngine = WSSecurityEngine.getInstance();
    
    
    public Wss4jSecurityInterceptorJWeb()
    {
        timestampStrict = true;
        validationTimeToLive = 300;
        securementTimeToLive = 300;
    }

    public void setSecurementActions(String securementActions)
    {
        this.securementActions = securementActions;
        securementActionsVector = new Vector();
        try
        {
            securementAction = WSSecurityUtil.decodeAction(securementActions, securementActionsVector);
        }
        catch(WSSecurityException ex)
        {
            throw new IllegalArgumentException(ex);
        }
    }

    public void setSecurementActor(String securementActor)
    {
        handler.setOption("actor", securementActor);
    }

    public void setSecurementCallbackHandler(CallbackHandler securementCallbackHandler)
    {
        handler.setSecurementCallbackHandler(securementCallbackHandler);
    }

    public void setSecurementCallbackHandlers(CallbackHandler securementCallbackHandler[])
    {
        handler.setSecurementCallbackHandler(new CallbackHandlerChain(securementCallbackHandler));
    }

    public void setSecurementEncryptionCrypto(Crypto securementEncryptionCrypto)
    {
        handler.setSecurementEncryptionCrypto(securementEncryptionCrypto);
    }

    public void setSecurementEncryptionEmbeddedKeyName(String securementEncryptionEmbeddedKeyName)
    {
        handler.setOption("EmbeddedKeyName", securementEncryptionEmbeddedKeyName);
    }

    public void setSecurementEncryptionKeyIdentifier(String securementEncryptionKeyIdentifier)
    {
        handler.setOption("encryptionKeyIdentifier", securementEncryptionKeyIdentifier);
    }

    public void setSecurementEncryptionKeyTransportAlgorithm(String securementEncryptionKeyTransportAlgorithm)
    {
        handler.setOption("encryptionKeyTransportAlgorithm", securementEncryptionKeyTransportAlgorithm);
    }

    public void setSecurementEncryptionParts(String securementEncryptionParts)
    {
        handler.setOption("encryptionParts", securementEncryptionParts);
    }

    public void setSecurementEncryptionSymAlgorithm(String securementEncryptionSymAlgorithm)
    {
        handler.setOption("encryptionSymAlgorithm", securementEncryptionSymAlgorithm);
    }

    public void setSecurementEncryptionUser(String securementEncryptionUser)
    {
        handler.setOption("encryptionUser", securementEncryptionUser);
    }

    public void setSecurementPassword(String securementPassword)
    {
        handler.setSecurementPassword(securementPassword);
    }

    public void setSecurementPasswordType(String securementUsernameTokenPasswordType)
    {
        handler.setOption("passwordType", securementUsernameTokenPasswordType);
    }

    public void setSecurementSignatureAlgorithm(String securementSignatureAlgorithm)
    {
        handler.setOption("signatureAlgorithm", securementSignatureAlgorithm);
    }

    public void setSecurementSignatureCrypto(Crypto securementSignatureCrypto)
    {
        handler.setSecurementSignatureCrypto(securementSignatureCrypto);
    }

    public void setSecurementSignatureKeyIdentifier(String securementSignatureKeyIdentifier)
    {
        handler.setOption("signatureKeyIdentifier", securementSignatureKeyIdentifier);
    }

    public void setSecurementSignatureParts(String securementSignatureParts)
    {
        handler.setOption("signatureParts", securementSignatureParts);
    }

    public void setSecurementSignatureUser(String securementSignatureUser)
    {
        handler.setOption("signatureUser", securementSignatureUser);
    }

    public void setSecurementUsername(String securementUsername)
    {
        this.securementUsername = securementUsername;
    }

    public void setSecurementTimeToLive(int securementTimeToLive)
    {
        if(securementTimeToLive <= 0)
        {
            throw new IllegalArgumentException("timeToLive must be positive");
        } else
        {
            this.securementTimeToLive = securementTimeToLive;
            return;
        }
    }

    public void setValidationTimeToLive(int validationTimeToLive)
    {
        if(validationTimeToLive <= 0)
        {
            throw new IllegalArgumentException("timeToLive must be positive");
        } else
        {
            this.validationTimeToLive = validationTimeToLive;
            return;
        }
    }

    /**
     * @deprecated Method setTimeToLive is deprecated
     */

    public void setTimeToLive(int timeToLive)
    {
        setValidationTimeToLive(timeToLive);
    }

    public void setValidationActions(String actions)
    {
        validationActions = actions;
        try
        {
            validationActionsVector = new Vector();
            validationAction = WSSecurityUtil.decodeAction(actions, validationActionsVector);
        }
        catch(WSSecurityException ex)
        {
            throw new IllegalArgumentException(ex);
        }
    }

    public void setValidationActor(String validationActor)
    {
        this.validationActor = validationActor;
    }

    public void setValidationCallbackHandler(CallbackHandler callbackHandler)
    {
        validationCallbackHandler = callbackHandler;
    }

    public void setValidationCallbackHandlers(CallbackHandler callbackHandler[])
    {
        validationCallbackHandler = new CallbackHandlerChain(callbackHandler);
    }

    public void setValidationDecryptionCrypto(Crypto decryptionCrypto)
    {
        validationDecryptionCrypto = decryptionCrypto;
    }

    public void setValidationSignatureCrypto(Crypto signatureCrypto)
    {
        validationSignatureCrypto = signatureCrypto;
    }

    public void setEnableSignatureConfirmation(boolean enableSignatureConfirmation)
    {
        handler.setOption("enableSignatureConfirmation", enableSignatureConfirmation);
        this.enableSignatureConfirmation = enableSignatureConfirmation;
    }

    public void setTimestampPrecisionInMilliseconds(boolean timestampPrecisionInMilliseconds)
    {
        handler.setOption("precisionInMilliseconds", timestampPrecisionInMilliseconds);
    }

    public void setTimestampStrict(boolean timestampStrict)
    {
        this.timestampStrict = timestampStrict;
    }

    public void setSecurementMustUnderstand(boolean securementMustUnderstand)
    {
        handler.setOption("mustUnderstand", securementMustUnderstand);
    }

    public void setSecurementUsernameTokenElements(String securementUsernameTokenElements)
    {
        handler.setOption("addUTElements", securementUsernameTokenElements);
    }

    public void afterPropertiesSet()
        throws Exception
    {
        Assert.isTrue(validationActions != null || securementActions != null, "validationActions or securementActions are required");
        if(validationActions != null)
        {
            if((validationAction & 1) != 0)
                Assert.notNull(validationCallbackHandler, "validationCallbackHandler is required");
            if((validationAction & 2) != 0)
                Assert.notNull(validationSignatureCrypto, "validationSignatureCrypto is required");
        }
        securityEngine.getWssConfig().setAllowNamespaceQualifiedPasswordTypes(true);
    }
    
    
    protected void secureMessage(SoapMessage soapMessage , MessageContext messageContext) throws WsSecuritySecurementException {
    	if(securementAction == 0 && !enableSignatureConfirmation){
    		return ;
    	}
    	if(logger.isDebugEnabled()){
    		logger.debug((new StringBuilder()).append("Securing message [").append(soapMessage).append("] with actions [").append(securementActions).append("]").toString());
    	}
    	RequestData requestData = initializeRequestData(messageContext);
    	Document envelopeAsDocument = toDocument(soapMessage , messageContext);
    	try {
			if(securementAction == 0){
				securementActionsVector = new Vector(0);
			}
			handler.doSenderAction(securementAction, envelopeAsDocument, requestData, securementActionsVector, false);
		} catch (WSSecurityException e) {
			throw new Wss4jSecuritySecurementException(e.getMessage() , e);
		}
    	
    	replaceMessage(soapMessage , envelopeAsDocument);
    }
    
    
    private RequestData initializeRequestData(MessageContext messageContext){
    	RequestData requestData = new RequestData();
    	requestData.setMsgContext(messageContext);
    	String contextUsername = (String)messageContext.getProperty("Wss4jSecurityInterceptor.securementUser");
    	if(StringUtils.hasLength(contextUsername)){
    		requestData.setUsername(contextUsername);
    	}
    	else {
    		requestData.setUsername(securementUsername);
    	}
    	requestData.setTimeToLive(securementTimeToLive);
    	return requestData ;
    }
    
    
    protected void validateMessage(SoapMessage soapMessage , MessageContext messageContext) throws WsSecurityValidationException {
    	if(logger.isDebugEnabled()){
    		logger.debug((new StringBuilder()).append("Validating message [").append(soapMessage).append("] with actions [").append(validationActions).append("]").toString());
    	}
    	if(validationAction == 0){
    		return ;
    	}
    	Document envelopeAsDocument = toDocument(soapMessage , messageContext);
    	try {
			Vector results = securityEngine.processSecurityHeader(envelopeAsDocument, validationActor, validationCallbackHandler, validationSignatureCrypto, validationDecryptionCrypto);
			if(CollectionUtils.isEmpty(results)){
				throw new Wss4jSecurityValidationException("No WS-Security header found");
			}
			checkResults(results , validationActionsVector);
			updateContextWithResults(messageContext , results);
    		verifyCertificateTrust(results);
    		verifyTimestamp(results);
    		processPrincipal(results);
		} catch (WSSecurityException ex) {
			throw new Wss4jSecurityValidationException(ex.getMessage(), ex);
		}
    	
		replaceMessage(soapMessage , envelopeAsDocument);
		soapMessage.getEnvelope().getHeader().removeHeaderElement(WS_SECURITY_NAME);
    }
    
    
    protected void checkResults(Vector results , Vector validationActionsVector) throws Wss4jSecurityValidationException {
    	if(!handler.checkReceiverResultsAnyOrder(results, validationActionsVector)){
    		throw new Wss4jSecurityValidationException("Security processing failed (actions mismatch)");
    	}
    	else {
    		return ;
    	}
    }
    
    
    private void updateContextWithResults(MessageContext messageContext , Vector results){
    	Vector handlerResults ;
    	if((handlerResults = (Vector)messageContext.getProperty("RECV_RESULTS")) == null){
    		handlerResults = new Vector();
    		messageContext.setProperty("RECV_RESULTS", handlerResults);
    	}
    	WSHandlerResult rResult = new WSHandlerResult(validationActor , results);
    	handlerResults.add(0 , rResult);
    	messageContext.setProperty("RECV_RESULTS", handlerResults);
    }
    
    
    protected void verifyCertificateTrust(Vector results) throws WSSecurityException {
    	RequestData requestData = new RequestData();
    	requestData.setSigCrypto(validationSignatureCrypto);
    	WSSecurityEngineResult actionResult = WSSecurityUtil.fetchActionResult(results,  2);
    	if(actionResult != null){
    		X509Certificate returnCert = (X509Certificate)actionResult.get("x509-certificate");
    		if(!handler.verifyTrust(returnCert,  requestData)){
    			throw new Wss4jSecurityValidationException("The certificate used for the signature is not trusted");
    		}
    	}
    }
    
    
	protected void verifyTimestamp(Vector results) throws WSSecurityException {
		WSSecurityEngineResult actionResult = WSSecurityUtil.fetchActionResult(results, 32);
		if (actionResult != null) {
			Timestamp timestamp = (Timestamp) actionResult.get("timestamp");
			if (timestamp != null
					&& timestampStrict
					&& !handler.verifyTimestamp(timestamp, validationTimeToLive))
				throw new Wss4jSecurityValidationException(
						(new StringBuilder()).append("Invalid timestamp : ")
								.append(timestamp.getID()).toString());
		}
	}
    
    
	private void processPrincipal(Vector results)
    {
        WSSecurityEngineResult actionResult = WSSecurityUtil.fetchActionResult(results, 1);
        if(actionResult != null)
        {
            Principal principal = (Principal)actionResult.get("principal");
            if(principal != null && (principal instanceof WSUsernameTokenPrincipal))
            {
                WSUsernameTokenPrincipal usernameTokenPrincipal = (WSUsernameTokenPrincipal)principal;
                UsernameTokenPrincipalCallback callback = new UsernameTokenPrincipalCallback(usernameTokenPrincipal);
                try
                {
                    validationCallbackHandler.handle(new Callback[] {
                        callback
                    });
                }
                catch(IOException ex)
                {
                    logger.warn("Principal callback resulted in IOException", ex);
                }
                catch(UnsupportedCallbackException ex) { }
            }
        }
    }
    
    
    protected Document toDocument(SoapMessage soapMessage , MessageContext messageContext) {
    	
    	if(soapMessage instanceof SaajSoapMessage){
    		throw new IllegalArgumentException("Message type not supported SaajSoapMessage");
    	}
    	
    	if(soapMessage instanceof org.springframework.ws.soap.axiom.AxiomSoapMessage){
    		return AxiomUtils.toDocument(((org.springframework.ws.soap.axiom.AxiomSoapMessage)soapMessage).getAxiomMessage().getSOAPEnvelope());
    	}
    	
    	if(soapMessage instanceof com.jweb.framework.ws.soap.axiom.AxiomSoapMessage){
    		return AxiomUtils.toDocument(((com.jweb.framework.ws.soap.axiom.AxiomSoapMessage)soapMessage).getAxiomMessage().getSOAPEnvelope());
    	}
    	
    	throw new IllegalArgumentException("Message type not supported [" + soapMessage + "]");
    }
    
    
    
    protected void replaceMessage(SoapMessage soapMessage , Document document){
    	
    	String soapAction ;
    	SOAPEnvelope soapEnvelope ;
    	SOAPFactory soapFactory ;
    	SOAPMessage localSoapMessage ;
    	
    	if(soapMessage instanceof org.springframework.ws.soap.axiom.AxiomSoapMessage){
    		soapAction = ((org.springframework.ws.soap.axiom.AxiomSoapMessage)soapMessage).getSoapAction();
    		soapEnvelope = AxiomUtils.toEnvelope(document);
    		soapFactory = (SOAPFactory)((org.springframework.ws.soap.axiom.AxiomSoapMessage)soapMessage).getAxiomMessage().getOMFactory();
    		localSoapMessage = soapFactory.createSOAPMessage();
    		localSoapMessage.setSOAPEnvelope(soapEnvelope);
    		
    		((org.springframework.ws.soap.axiom.AxiomSoapMessage)soapMessage).setAxiomMessage(localSoapMessage);
    		((org.springframework.ws.soap.axiom.AxiomSoapMessage)soapMessage).setSoapAction(soapAction);
    	}
    	else if(soapMessage instanceof com.jweb.framework.ws.soap.axiom.AxiomSoapMessage){
    		
    		soapAction = ((com.jweb.framework.ws.soap.axiom.AxiomSoapMessage)soapMessage).getSoapAction();
    		soapEnvelope = AxiomUtils.toEnvelope(document);
    		soapFactory = (SOAPFactory)((com.jweb.framework.ws.soap.axiom.AxiomSoapMessage)soapMessage).getAxiomMessage().getOMFactory();
    		localSoapMessage = soapFactory.createSOAPMessage();
    		localSoapMessage.setSOAPEnvelope(soapEnvelope);
    		
    		((com.jweb.framework.ws.soap.axiom.AxiomSoapMessage)soapMessage).setAxiomMessage(localSoapMessage);
    		((com.jweb.framework.ws.soap.axiom.AxiomSoapMessage)soapMessage).setSoapAction(soapAction);
    	}
    }
    
    
    protected void cleanUp()
    {
        if(validationCallbackHandler != null)
            try
            {
                CleanupCallback cleanupCallback = new CleanupCallback();
                validationCallbackHandler.handle(new Callback[] {
                    cleanupCallback
                });
            }
            catch(IOException ex)
            {
                logger.warn("Cleanup callback resulted in IOException", ex);
            }
            catch(UnsupportedCallbackException ex) { }
    }
    
    
	
}
